import { BaseEmbeddings } from '@cherrystudio/embedjs-interfaces'
import { VoyageEmbeddings as _VoyageEmbeddings } from '@langchain/community/embeddings/voyage'

/**
 * VoyageEmbeddings 封装类
 * 继承 BaseEmbeddings，统一对外提供文本向量化接口
 */
export default class VoyageEmbeddings extends BaseEmbeddings {
  // VoyageEmbeddings 模型实例
  private model: _VoyageEmbeddings

  /**
   * 构造函数
   * @param configuration 配置参数（模型名、输出维度等）
   */
  constructor(private readonly configuration?: ConstructorParameters<typeof _VoyageEmbeddings>[0]) {
    super()
    // 若未传入配置则初始化为空对象
    if (!this.configuration) this.configuration = {}
    // 若未指定模型名，默认使用 'voyage-3'
    if (!this.configuration.modelName) this.configuration.modelName = 'voyage-3'

    // 若未指定输出维度则抛出异常
    if (!this.configuration.outputDimension) {
      throw new Error('You need to pass in the optional dimensions parameter for this model')
    }
    // 初始化 VoyageEmbeddings 实例
    this.model = new _VoyageEmbeddings(this.configuration)
  }

  /**
   * 获取向量维度
   * @returns 当前模型的输出维度
   */
  override async getDimensions(): Promise<number> {
    if (!this.configuration?.outputDimension) {
      throw new Error('You need to pass in the optional dimensions parameter for this model')
    }
    return this.configuration?.outputDimension
  }

  /**
   * 批量文本向量化
   * @param texts 文本数组
   * @returns 二维向量数组
   */
  override async embedDocuments(texts: string[]): Promise<number[][]> {
    return this.model.embedDocuments(texts)
  }

  /**
   * 单条文本向量化
   * @param text 单条文本
   * @returns 一维向量数组
   */
  override async embedQuery(text: string): Promise<number[]> {
    return this.model.embedQuery(text)
  }
}
